import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F

def check_device():
    """检查并返回可用设备"""
    if torch.cuda.is_available():
        device = torch.device("cuda")
        print(f"使用GPU设备: {torch.cuda.get_device_name(0)}")
    else:
        device = torch.device("cpu")
        print("使用CPU设备")
    return device

# 定义一个简单的神经网络
class SimpleNet(nn.Module):
    def __init__(self, input_size=2, hidden_size=3, output_size=1):
        super(SimpleNet, self).__init__()
        # 定义网络层
        self.fc1 = nn.Linear(input_size, hidden_size)  # 输入层到隐藏层
        self.fc2 = nn.Linear(hidden_size, output_size)  # 隐藏层到输出层
    
    def forward(self, x):
        # 前向传播
        x = F.relu(self.fc1(x))  # 隐藏层使用ReLU激活函数
        x = torch.sigmoid(self.fc2(x))  # 输出层使用Sigmoid激活函数
        return x

def create_sample_data(device):
    """创建示例数据并移到指定设备"""
    # 创建一些示例数据 (XOR问题)
    X = torch.tensor([[0.0, 0.0],
                      [0.0, 1.0],
                      [1.0, 0.0],
                      [1.0, 1.0]], dtype=torch.float32)
    
    # XOR问题的目标输出
    y = torch.tensor([[0.0],
                      [1.0],
                      [1.0],
                      [0.0]], dtype=torch.float32)
    
    # 将数据移到指定设备
    X = X.to(device)
    y = y.to(device)
    
    return X, y

def train_network():
    """训练神经网络"""
    # 检查设备
    device = check_device()
    
    # 创建网络实例并移到设备
    model = SimpleNet().to(device)
    print("网络结构:")
    print(model)
    
    # 打印模型参数名称和形状
    print("\n=== 模型参数 ===")
    for name, param in model.named_parameters():
        print(f"参数名称: {name}, 形状: {param.shape}, 设备: {param.device}")
    
    # 创建损失函数和优化器
    criterion = nn.MSELoss()  # 均方误差损失
    optimizer = optim.SGD(model.parameters(), lr=0.1)  # 随机梯度下降优化器
    
    # 创建示例数据并确保在相同设备上
    X, y = create_sample_data(device)
    print(f"\n数据设备: X在{X.device}, y在{y.device}")
    
    # 验证模型和数据是否在同一设备
    model_device = next(model.parameters()).device
    print(f"模型设备: {model_device}")
    print(f"设备一致性检查: {model_device == X.device}")
    
    # 训练前向传播示例
    print("\n=== 训练前的前向传播 ===")
    with torch.no_grad():  # 不计算梯度
        output_before = model(X)
        loss_before = criterion(output_before, y)
        print(f"输入数据:\n{X.cpu().numpy()}")
        print(f"目标输出:\n{y.squeeze().cpu().numpy()}")
        print(f"训练前输出:\n{output_before.squeeze().cpu().numpy()}")
        print(f"训练前损失: {loss_before.item():.6f}")
    
    # 训练网络
    print("\n=== 开始训练 ===")
    epochs = 1000
    losses = []
    
    for epoch in range(epochs):
        # 前向传播
        outputs = model(X)
        loss = criterion(outputs, y)
        
        # 反向传播和优化
        optimizer.zero_grad()  # 清零梯度
        loss.backward()        # 计算梯度
        optimizer.step()       # 更新权重
        
        losses.append(loss.item())
        
        # 每100个epoch打印一次
        if (epoch + 1) % 100 == 0:
            print(f'Epoch [{epoch+1}/{epochs}], Loss: {loss.item():.6f}')
    
    # 训练后前向传播示例
    print("\n=== 训练后的前向传播 ===")
    with torch.no_grad():  # 不计算梯度
        output_after = model(X)
        loss_after = criterion(output_after, y)
        print(f"训练后输出:\n{output_after.squeeze().cpu().numpy()}")
        print(f"训练后损失: {loss_after.item():.6f}")
    
    return model, X, y, losses

def demonstrate_autograd(device):
    """演示自动求导机制"""
    print("\n=== 自动求导机制演示 ===")
    
    # 创建需要计算梯度的张量并移到设备
    x = torch.tensor([2.0, 3.0], requires_grad=True).to(device)
    print(f"输入 x: {x} (在 {x.device} 上)")
    
    # 定义计算过程
    y = x ** 2
    print(f"计算 y = x^2: {y}")
    
    z = y.sum()
    print(f"计算 z = sum(y): {z}")
    
    # 计算梯度
    z.backward()
    
    # 打印梯度
    print(f"x 的梯度 (dz/dx): {x.grad}")
    
    # 数学验证: 
    # z = x[0]^2 + x[1]^2
    # dz/dx[0] = 2*x[0] = 2*2 = 4
    # dz/dx[1] = 2*x[1] = 2*3 = 6
    print("数学验证:")
    print(f"dz/dx[0] = 2*x[0] = 2*{x[0].item()} = {2*x[0].item()}")
    print(f"dz/dx[1] = 2*x[1] = 2*{x[1].item()} = {2*x[1].item()}")

def print_computation_graph_info(model, X):
    """打印计算图相关信息"""
    print("\n=== 计算图信息 ===")
    print("PyTorch使用动态计算图（Define-by-Run）:")
    print("1. 计算图在前向传播过程中动态构建")
    print("2. 每次前向传播都会创建新的计算图")
    print("3. 反向传播后计算图会被释放")
    print("4. 通过requires_grad=True的张量可以追踪梯度")
    
    # 创建一个简单的示例来展示计算图
    x = torch.tensor([[0.5, 0.8]], dtype=torch.float32, requires_grad=True).to(X.device)
    print(f"\n示例输入: {x} (在 {x.device} 上)")
    
    # 前向传播
    y = model(x)
    print(f"前向传播输出: {y} (在 {y.device} 上)")
    
    # 反向传播
    y.backward()
    print("已完成反向传播")
    
    # 显示梯度信息
    print("\n梯度信息:")
    for name, param in model.named_parameters():
        if param.grad is not None:
            print(f"{name} 的梯度: {param.grad} (在 {param.grad.device} 上)")

def describe_network_architecture():
    """描述网络架构"""
    print("\n=== 神经网络架构说明 ===")
    print("网络结构: 2-3-1 (输入层-隐藏层-输出层)")
    print("输入层: 2个节点 (处理二维输入数据)")
    print("隐藏层: 3个节点，使用ReLU激活函数")
    print("输出层: 1个节点，使用Sigmoid激活函数")
    print("\n权重矩阵形状:")
    print("  W1 (输入→隐藏): [2, 3] - 2个输入连接到3个隐藏节点")
    print("  W2 (隐藏→输出): [3, 1] - 3个隐藏节点连接到1个输出节点")
    print("\n偏置向量形状:")
    print("  b1 (隐藏层): [1, 3] - 3个隐藏节点的偏置")
    print("  b2 (输出层): [1, 1] - 1个输出节点的偏置")

def explain_device_placement():
    """解释设备放置的重要性"""
    print("\n=== 设备放置说明 ===")
    print("在PyTorch中，模型和数据必须在相同的设备上才能进行计算:")
    print("1. CPU设备: 所有计算在CPU上进行")
    print("2. GPU设备: 所有计算在GPU上进行，速度更快")
    print("\n关键操作:")
    print("- model.to(device): 将模型移到指定设备")
    print("- tensor.to(device): 将张量移到指定设备")
    print("- device = torch.device('cuda' if torch.cuda.is_available() else 'cpu'): 检查并选择设备")
    print("\n常见错误:")
    print("- 模型在GPU上，数据在CPU上 → RuntimeError")
    print("- 模型在CPU上，数据在GPU上 → RuntimeError")
    print("- 解决方法: 确保模型和数据都在同一设备上")

def main():
    """主函数"""
    print("PyTorch神经网络设备兼容性示例")
    print("=" * 50)
    
    # 描述网络架构
    describe_network_architecture()
    
    # 解释设备放置
    explain_device_placement()
    
    # 演示自动求导机制
    device = check_device()
    demonstrate_autograd(device)
    
    # 训练网络
    model, X, y, losses = train_network()
    
    # 打印计算图信息
    print_computation_graph_info(model, X)
    
    # 显示最终结果
    print("\n=== 训练总结 ===")
    print(f"最终损失: {losses[-1]:.6f}")
    print("通过训练，网络学会了XOR函数的映射关系")
    print("这证明了神经网络可以学习非线性函数")
    
    print("\n程序执行完成!")

if __name__ == '__main__':
    main()